# 从tuping模块导入Literal,用于类型注解，限制变量的取值范围
import base64
import os
from io import BytesIO
# 导入streamlit 模块 用于创建web应用
from typing import Literal

# 导入streamlit 模块 用于创建web应用
import streamlit as st
from langchain_community.chat_models import ChatOllama
from langchain_openai import ChatOpenAI
from streamlit_flow import StreamlitFlowState, StreamlitFlowNode, StreamlitFlowEdge, streamlit_flow
from streamlit_flow.layouts import TreeLayout
from xinference_client.client.restful.restful_client import Client

# 定义支持的平台名称
PLATFORMS = ["xunfei","modelscope", "openAI", "Ollama", "Xinference"]


# 获取指定平台的模型列表
def get_llm_models(platform_type: Literal[tuple[PLATFORMS]], base_url: str = "", api_key: str = ""):
    # 如果平台类型是Ollama
    ### Ollama 像一个开箱即用的“模型播放器”。它的核心目标是让用户在个人电脑上以最简单、最无痛的方式下载和运行各种开源大模型（尤其是 Llama、Mistral 系列）。用户体验极佳，命令行操作非常简单。
    if platform_type == "Ollama":
        try:
            # 导入 ollama 模块
            import ollama
            # 如果没有提供base_url ，使用默认的本地地址
            if not base_url:
                base_url = "http://127.0.0.1:11434"
            # 创建 Ollama 客户端
            client = ollama.Client(host=base_url)
            # 获取模型列表，排除包含“bert”的模型
            # llm_models = [model["model"] for model in client.list()["models"] if "bert" not in model.details.families]
            llm_models = [
                model["model"]
                for model in client.list()["models"]
                if not any("bert" in fam for fam in model.details.families)
            ]
            # 返回模型列表
            return llm_models
        except Exception as e:
            # 如果发生异常，显示错误信息 并返回空列表
            st.toast(f"尝试链接{platform_type} 获取LLM模型时发生错误,\n{e}")
    ### Xinference 更像一个可扩展的“模型服务引擎”。它的核心目标是让用户在任何环境（从笔记本电脑到百台服务器的集群）中部署和管理模型，并将其以标准的 API 形式提供服务，方便集成到其他应用中。它更注重生产环境的部署和管理能力
    elif platform_type == "Xinference":
        try:
            # 从xinference_client 模块导入 RESTfulClient类并重命名为Client
            from xinference_client import RESTfulClient as client
            # 如果没有提供base_url,使用默认本地地址
            if not base_url:
                base_url = "http://127.0.0.1:9997"
            # 创建xinference 客户端
            client = Client(base_url=base_url)
            # 获取模型列表，筛选出类型为LLM的模型
            llm_models = client.list_models()
            return [k for k, v in llm_models.items() if v.get("model_type") == "LLM"]
        except Exception as e:
            # 如果发生异常，显示错误信息并返回空列表
            st.toast(f"尝试连接{platform_type} 获取 LLM 模型时 发生错误：\n{e}")
            return []
    # 如果平台类型是ZhipuAI
    elif platform_type == "ZhipuAI":
        # from zhipuai import ZhipuAI
        # client = ZhipuAI(api_key="")
        # llm_models = client.list_models()
        # 返回一个预定义的模型列表
        return [
            'glm-4-alltools',
            'glm-4-plus',
            'glm-4-0520',
            'glm-4',
            'glm-4-air',
            'glm-4-airx',
            'glm-4-long',
            'glm-4-flashx',
            'glm-4-flash'
        ]
    # 如果平台类型是OpenAI
    elif platform_type == "OpenAI":
        # 这里是注释掉的代码，表示如何从OpenAI获取模型列表
        # from zhipuai import ZhipuAI
        #
        # client = ZhipuAI(
        #     api_key="",  # 填写您的 APIKey
        # )
        # client.list_models()
        # 返回一个预定义的模型列表
        return [
            'gpt-4o',
            'gpt-3.5-turbo'
        ]
    # 如果平台类型是xunfei
    elif platform_type == "xunfei":
        ## 直接返回 本人配置了可用的讯飞模型 使用OpenAi sdk调用
        return [
            'xdeepseekv32exp',
            'xop3qwen1b7'
        ]
    # 如果平台类型是modelscope
    elif platform_type == "modelscope":
        ## 直接返回 本人配置了可用的讯飞模型 使用OpenAi sdk调用
        return [
            'Qwen/Qwen2.5-VL-7B-Instruct'
        ]


# 定义一个喊出 获取可用的向量模型
def get_embedding_models(platform_type: Literal[tuple[PLATFORMS]], base_url: str = "", api_key: str = ""):
    if platform_type == "Ollama":
        try:
            import ollama
            if not base_url:
                base_url = "http://127.0.0.1:11434"
            client = ollama.Client(host=base_url)
            # 获取模型列表，筛选出 包含 bert的模型
            # embedding_models = [model["model"] for model in client.list()["models"] if "bert" in model.details.families]
            embedding_models = [
                model["model"]
                for model in client.list()["models"]
                if any("bert" in fam for fam in model.details.families)
            ]
            # 返回模型列表
            return embedding_models
        except Exception as e:
            # 如果发生异常，显示错误信息并返回空列表
            st.toast(f"尝试连接 {platform_type} 获取 Embedding 模型时发生错误：\n{e}")
            return []
        # 如果平台类型是Xinference
    elif platform_type == "Xinference":
        try:
            # 从xinference_client模块导入RESTfulClient类并重命名为Client
            from xinference_client import RESTfulClient as Client
            # 如果没有提供base_url，使用默认的本地地址
            if not base_url:
                base_url = "http://127.0.0.1:9997"
            # 创建Xinference客户端
            client = Client(base_url=base_url)
            # 获取模型列表，筛选出类型为embedding的模型
            embedding_models = client.list_models()
            return [k for k, v in embedding_models.items() if v.get("model_type") == "embedding"]
        except Exception as e:
            # 如果发生异常，显示错误信息并返回空列表
            st.toast(f"尝试连接 {platform_type} 获取 Embedding 模型时发生错误：\n{e}")
            return []


# 定义一个函数获取 指定平台的聊天模型
def get_chatllm(
        platform_type: Literal[tuple[PLATFORMS]],
        model: str,
        base_url: str = "",
        api_key: str = "",
        temperature: float = 0.1
):
    # 如果平台式 ollama
    if platform_type == "Ollama":
        # 如果没有指定 base_url，默认使用本地地址
        if not base_url:
            base_url = "http://127.0.0.1:11434"
        # ollama 模型使用 ollamasdk 下面几款使用通用openaisdk
        return ChatOllama(
            temperature=temperature,
            # streaming = True, # 是否启用流式传输
            model=model,
            base_url=base_url,
        )
    # 如果平台是 Xinference
    elif platform_type == "Xinference":
        if not base_url:
            base_url = "http://127.0.0.1:9997/v1"
        # 如果没有提供 api_key 使用默认的EMPTY
        if not api_key:
            api_key = "EMPTY"
    # 如果平台类型是ZhipuAI
    elif platform_type == "ZhipuAI":
        if not base_url:
            base_url = "https://api.bigmodel.cn/api/pass/v4"
        if not api_key:
            api_key = "EMPTY"
    # 如果平台类型是 OpenAI
    elif platform_type == "OpenAI":
        if not base_url:
            base_url = "http://api/openai.com/v1"
        if not api_key:
            api_key = os.getenv("OPENAI_API_KEY")
    elif platform_type == "xunfei":
        if not base_url:
            base_url = "https://maas-api.cn-huabei-1.xf-yun.com/v1"
        if not api_key:
            api_key = os.getenv("XUNFEI_OPENAI_API_KEY")
    elif platform_type == "modelscope":
        if not base_url:
            base_url = "https://api-inference.modelscope.cn/v1"
        if not api_key:
            api_key = os.getenv("MODELSCOPE_OPENAI_API_KEY")

    # 返回一个ChatOpenAI对象，用于与OpenAI的聊天模型交互
    return ChatOpenAI(
        temperature=temperature,
        model_name=model,
        streaming=True,
        base_url=base_url,
        api_key=api_key,
    )


# 定义一个函数show_graph，用于在Streamlit中展示流程图
def show_graph(graph):
    # 创建一个streamlitFlowState 对象，包含节点和边的信息
    # pip install streamlit-flow-component
    flow_state = StreamlitFlowState(
        # 创建节点列表 ，每个节点包含id、位置、数据和类型
        nodes=[
            StreamlitFlowNode(
                id=node.id,
                pos=(0, 0),
                data={"content": node.id},
                node_type="input" if node.id == "__start__"
                else "output" if node.id == "__end__"
                else "default",
            ) for node in graph.nodes.values()
        ],
        #  创建边列表，每条边包含id、源节点，目标节点和动画效果
        edges=[StreamlitFlowEdge(
            id=str(enum),
            source=edge.source,
            target=edge.target,
            animate=True
        ) for enum, edge in enumerate(graph.edges)]
    )
    # 使用streamlit_flow 函数展示流程图，指定布局和视图设置
    streamlit_flow(
        "example_flow",
        flow_state,
        layout=TreeLayout(direction="down"),
        fit_view=True
    )


# 定义一个函数 get_kb_names获取知识库名称列表
def get_kb_names():
    # 获取当前文件所在目录下的kb目录路径
    kb_root = os.path.join(os.path.dirname(__file__), "kb")
    # 如果kb目录不存在 则创建该目录
    if not os.path.exists(kb_root):
        os.mkdir(kb_root)
    # 获取kb目录下的所有子目录名称，并返回这些名称
    kb_names = [f for f in os.listdir(kb_root) if os.path.isdir(os.path.join(kb_root, f))]
    return kb_names


# 定义一个函数获取向量模型
def get_embedding_model(
        platform_type: Literal[tuple[PLATFORMS]] = "Ollama",
        model: str = "text-embedding-ada-002",
        base_url: str = os.getenv('OPENAI_BASE_URL'),
        api_key: str = os.getenv('OPENAI_API_KEY')
):
    if platform_type == "Ollama":
        # pip install langchain-ollama
        from langchain_ollama import OllamaEmbeddings
        return OllamaEmbeddings(base_url=base_url, model=model)
    elif platform_type == "Xinference":
        # 从langchain_community.embeddings.xinference模块导入XinferenceEmbeddings类
        from langchain_community.embeddings.xinference import XinferenceEmbeddings
        # 如果没有提供base_url，使用默认的本地地址
        if not base_url:
            base_url = "http://127.0.0.1:9997/v1"
        # 返回一个XinferenceEmbeddings对象，用于与Xinference的Embedding模型交互
        return XinferenceEmbeddings(server_url=base_url, model_uid=model)
    # 否则，使用OpenAI的Embedding模型
    else:
        # 从langchain_openai模块导入OpenAIEmbeddings类
        from langchain_openai import OpenAIEmbeddings
        # 返回一个OpenAIEmbeddings对象，用于与OpenAI的Embedding模型交互
        return OpenAIEmbeddings(base_url=base_url, api_key=api_key, model=model)


# 定义函数 获取图片base64 编码
def get_img_base64(file_name: str) -> str:
    # 获取当前文件所在目录下的img目录中的图片路径
    image_path = os.path.join(os.path.dirname(__file__), "img", file_name)
    # 获取图片
    with open(image_path, "rb") as f:
        # 将图片数据读入内存
        buffer = BytesIO(f.read())
        # 对图片进行base64 编码 ，并解码为字符串
        base_str = base64.b64encode(buffer.getvalue()).decode("utf-8")

    # 返回 base64编码的图片数据，格式为 data:image/png;base64,编码数据
    return f"data:image/jpeg;base64,{base_str}"


if __name__ == "__main__":
    # 打印Ollama 平台idea默认Embedding模型
    print(get_embedding_models(platform_type="Ollama"))
    # print(get_llm_models(platform_type="Ollama"))
